"""可视化模块 - 用于在2D地图上显示卫星轨迹和访问窗口"""

import numpy as np
from matplotlib import pyplot as plt
import plotly.graph_objects as go
from astropy import units as u
from astropy.time import Time, TimeDelta
from astropy.coordinates import GCRS, ITRS, CartesianRepresentation, SphericalRepresentation
from poliastro.bodies import Earth
from poliastro.earth.plotting.groundtrack import GroundtrackPlotter
from poliastro.twobody.sampling import EpochsArray
from poliastro.earth import EarthSatellite
from poliastro.spacecraft import Spacecraft
from poliastro.frames import Planes
from typing import List

from target import GroundTarget, GroundTargetGroup

class ConstellationVisualizer:
    """星座可视化器"""
    
    def __init__(self):
        """初始化星座可视化器"""
        self.fig = None
        self.groundtrack_plotter = None
        
    def initialize_plot(self):
        """初始化绘图"""
        # 使用plotly创建基本图形
        self.fig = go.Figure(go.Scattergeo())
        self.groundtrack_plotter = GroundtrackPlotter(self.fig)
        
        # 设置基本布局
        self.groundtrack_plotter.update_layout(
            title="卫星星座地面轨迹",
            height=800,
            width=1200,
            geo=dict(
                showland=True,
                showcountries=True,
                showocean=True,
                countrywidth=0.5,
                landcolor='rgb(230, 230, 230)',
                oceancolor='rgb(220, 240, 250)',
                projection=dict(
                    type='orthographic',
                    rotation=dict(lon=0, lat=0, roll=0)
                )
            )
        )
        
        return self.groundtrack_plotter.fig
        
    def plot_constellation_groundtracks(self, ground_targets:List[GroundTarget], access_results):
        """
        绘制星座地面轨迹
        
        参数:
            ground_targets: 地面观测(被观测)目标
            access_results: 预计算的访问窗口和卫星轨道数据
        """
        if self.groundtrack_plotter is None:
            self.initialize_plot()
        
        # 绘制地面目标
        for target in ground_targets:
            # 添加目标点标记，使用不同图标或形状
            target_trace = go.Scattergeo(
                lat=[target.latitude.to(u.deg).value],
                lon=[target.longitude.to(u.deg).value],
                mode="markers+text",
                name=f"目标: {target.name}",
                marker=dict(
                    size=10,
                    color="black",
                    symbol="star"  # 可以为不同目标设置不同的symbol
                ),
                text=[target.name],
                textposition="top center"
            )
            self.groundtrack_plotter.add_trace(target_trace)
        
        # 颜色列表，用于区分不同卫星
        colors = ["red", "blue", "green", "purple", "orange", "cyan", "magenta", "yellow"]
        
        # 绘制星座轨迹（低饱和度、线条较细或虚线）
        for i, (satellite, _) in enumerate(access_results):
            color = colors[i % len(colors)]
            color_desaturated = self._desaturate_color(color, 0.6)  # 降低颜色饱和度
            
            # 提取轨迹点数据
            orbit_data = satellite._precomputed_data
            times = orbit_data['times']
            # 从itrs计算经纬度
            lats = []
            lons = []
            
            for j, time in enumerate(times):
                itrs = orbit_data['itrs'][j]
                lat = itrs.represent_as(SphericalRepresentation).lat.to(u.deg).value
                lon = itrs.represent_as(SphericalRepresentation).lon.to(u.deg).value
                lats.append(lat)
                lons.append(lon)
            
            # 绘制轨迹（低饱和度）
            trace = go.Scattergeo(
                lat=lats,
                lon=lons,
                mode="lines",
                name=f"{satellite.name} 轨迹",
                line=dict(
                    color=color_desaturated,
                    width=1,
                    dash='dot'  # 虚线
                ),
                showlegend=True,
                hovertemplate="经度: %{lon:.2f}°, 纬度: %{lat:.2f}°, "+f"{satellite.name} 轨迹"+"<br>%{text}<extra></extra>",
                text=[f"时间: {time.isot}" for time in times],  # 添加时间戳信息
                visible='legendonly'  # 默认不显示该线段
            )
            self.groundtrack_plotter.add_trace(trace)
            
            # 绘制卫星起始点（实心圆点）
            if lats and lons:
                start_point = go.Scattergeo(
                    lat=[lats[0]],
                    lon=[lons[0]],
                    mode="markers",
                    name=f"{satellite.name} 起始点",
                    marker=dict(
                        size=8,
                        color=color,
                        symbol="circle"
                    ),
                    showlegend=True,
                    hovertemplate="经度: %{lon:.2f}°, 纬度: %{lat:.2f}°, "+f"{satellite.name} 起始点"+"<br>%{text}<extra></extra>",
                    text=[f"时间: {times[0].isot}"]  # 添加起始时间信息
                )
                self.groundtrack_plotter.add_trace(start_point)
        
        # 绘制访问窗口（高饱和度、明亮线条）
        if access_results:
            self._plot_access_windows(access_results)
        
        return self.groundtrack_plotter.fig
    
    def _desaturate_color(self, color_name, saturation_factor):
        """
        降低颜色饱和度
        
        参数:
            color_name: 颜色名称
            saturation_factor: 饱和度因子（0-1）
        
        返回:
            降低饱和度后的RGB颜色字符串
        """
        import matplotlib.colors as mcolors
        
        # 转换颜色名称为RGB
        rgb = mcolors.to_rgb(color_name)
        
        # 转换为HSV
        import colorsys
        h, s, v = colorsys.rgb_to_hsv(*rgb)
        
        # 降低饱和度
        s = s * saturation_factor
        
        # 转回RGB
        r, g, b = colorsys.hsv_to_rgb(h, s, v)
        
        # 返回RGB字符串
        return f"rgb({int(r*255)}, {int(g*255)}, {int(b*255)})"
    
    def _plot_access_windows(self, access_results):
        """
        绘制访问窗口
        
        参数:
            access_results: 访问窗口列表
        """
        # 颜色列表，用于区分不同卫星
        colors = ["red", "blue", "green", "purple", "orange", "cyan", "magenta", "yellow"]
        
        # 创建一个字典，用于跟踪已经处理过的访问窗口
        processed_windows = {}
        
        # 处理每个卫星的访问结果
        for i, (satellite, target_results) in enumerate(access_results):
            color = colors[i % len(colors)]
            color_desaturated = self._desaturate_color(color, 0.95)  # 降低颜色饱和度
            orbit_data = satellite._precomputed_data
            
            # 处理每个目标的访问窗口
            for target, windows in target_results:
                # 为每个窗口绘制高亮轨迹
                for window in windows:
                    # 生成唯一标识符，避免重复处理同一窗口
                    window_id = f"{satellite.name}_{target.name}_{window.start_time.isot}_{window.end_time.isot if window.end_time else 'ongoing'}"
                    if window_id in processed_windows:
                        continue
                    
                    processed_windows[window_id] = True
                    
                    # 找出窗口时间范围内的轨迹点
                    window_times = []
                    window_lats = []
                    window_lons = []
                    
                    # 从预计算轨道数据中提取窗口期间的点
                    for j, time in enumerate(orbit_data['times']):
                        if window.start_time <= time and (window.end_time is None or time <= window.end_time):
                            window_times.append(time)
                            itrs = orbit_data['itrs'][j]
                            lat = itrs.represent_as(SphericalRepresentation).lat.to(u.deg).value
                            lon = itrs.represent_as(SphericalRepresentation).lon.to(u.deg).value
                            window_lats.append(lat)
                            window_lons.append(lon)
                    
                    # 确定访问方向（卫星->地面目标 或 地面目标->卫星）
                    if hasattr(satellite, 'orbit') and hasattr(target, 'latitude'):
                        direction = f"{satellite.name} → {target.name}"  # 卫星->地面目标
                    elif hasattr(target, 'orbit') and hasattr(satellite, 'latitude'):
                        direction = f"{satellite.name} ← {target.name}"  # 地面目标->卫星
                    else:
                        direction = f"{satellite.name} ↔ {target.name}"  # 未知方向
                    
                    # 绘制高亮轨迹
                    if window_lats and window_lons:
                        trace = go.Scattergeo(
                            lat=window_lats,
                            lon=window_lons,
                            mode="lines",
                            name=direction,
                            line=dict(
                                color=color_desaturated,
                                width=2  # 加粗线条
                            ),
                            showlegend=True,
                            hovertemplate="经度: %{lon:.2f}°, 纬度: %{lat:.2f}°, "+direction+"<br>%{text}<extra></extra>",
                            text=[f"时间: {time.isot}" for time in window_times]  # 添加时间戳信息
                        )
                        self.groundtrack_plotter.add_trace(trace) 